from llm import  llm
from llm import super_eval
import json

import requests
from config import Config


def exec_sql_s(sql):
    headers = {
        "Authorization": f"Bearer {Config.team_key}",
        "Accept": "application/json"
    }
    url = "https://comm.chatglm.cn/finglm2/api/query"

    response = requests.post(url, headers=headers, json={
        "sql": sql,
        "limit": 100
    })
    try:
        if 'data' not in response.json():
            print(response.json())
        return response.json()['data']
    except:
        return []

def output_result(result, info_list):
    if 'data' in result and len(result['data']) > 0:
        info_list.append(json.dumps(result['data'], ensure_ascii=False, indent=1) + '\n')

    return info_list


def process_company_name(value,is_jj = False):
    res_lst = []
    # 处理公司名称，在三个表中搜索匹配的记录
    tables = ['ConstantDB.SecuMain', 'ConstantDB.HK_SecuMain', 'ConstantDB.US_SecuMain']
    columns_to_match = ['CompanyCode', 'SecuCode', 'ChiName', 'ChiNameAbbr',
                        'EngName', 'EngNameAbbr', 'SecuAbbr', 'ChiSpelling']
    columns_to_select = ['InnerCode', 'CompanyCode', 'SecuCode', 'ChiName', 'ChiNameAbbr',
                         'EngName', 'EngNameAbbr', 'SecuAbbr', 'ChiSpelling']

    value = value.replace("'", "''")  # 防止 SQL 注入

    for table in tables:
        if 'US' in table:
            columns_to_match.remove('ChiNameAbbr')
            columns_to_select.remove('ChiNameAbbr')
            columns_to_match.remove('EngNameAbbr')
            columns_to_select.remove('EngNameAbbr')

        match_conditions = [f"{col} = '{value}'" for col in columns_to_match]
        where_clause = ' OR '.join(match_conditions)
        sql = f"""
        SELECT {', '.join(columns_to_select)}
        FROM {table}
        WHERE {where_clause}
        """
        # print(sql)
        result = exec_sql_s(sql)
        if result:
            res_lst.append((result, table))
        else:
            continue
    else:
        if is_jj and not res_lst:
            return _process_jj(value)
        # print(f"未在任何表中找到公司名称为 {value} 的信息。")

    return res_lst


def process_code(value):
    res_lst = []
    # 处理代码，在三个表中搜索匹配的记录
    tables = ['ConstantDB.SecuMain', 'ConstantDB.HK_SecuMain', 'ConstantDB.US_SecuMain']
    columns_to_select = ['InnerCode', 'CompanyCode', 'SecuCode', 'ChiName', 'ChiNameAbbr',
                         'EngName', 'EngNameAbbr', 'SecuAbbr', 'ChiSpelling']

    value = value.replace("'", "''")  # 防止 SQL 注入

    for table in tables:
        if 'US' in table:
            columns_to_select.remove('ChiNameAbbr')
            columns_to_select.remove('EngNameAbbr')

        sql = f"""
        SELECT {', '.join(columns_to_select)}
        FROM {table}
        WHERE SecuCode = '{value}'
        """
        result = exec_sql_s(sql)
        if result:
            res_lst.append((result, table))
        else:
            continue
    else:
        ...
        # print(f"未在任何表中找到代码为 {value} 的信息。")
    return res_lst


def process_bk(value):
    sql = f"SELECT ClassName, ClassCode FROM AStockIndustryDB.LC_ConceptList WHERE ClassName = '{value}' LIMIT 1"
    res = exec_sql_s(sql)
    if res:
        return [(res, f'{value}是AStockIndustryDB.LC_ConceptList的一级概念')]
    else:
        sql = f"SELECT SubclassName, SubclassCode FROM AStockIndustryDB.LC_ConceptList WHERE SubclassName = '{value}' LIMIT 1"
        res = exec_sql_s(sql)
        if res:
            return [(res, f'{value}是AStockIndustryDB.LC_ConceptList的二级概念，其子概念是三级概念，请使用SubclassName或SubclassCode 进行接下来的查询')]
        else:
            sql = f"SELECT ConceptName, ConceptCode FROM AStockIndustryDB.LC_ConceptList WHERE ConceptName = '{value}' LIMIT 1"
            res = exec_sql_s(sql)
            if res:
                return [(res, f'{value}是AStockIndustryDB.LC_ConceptList的三级概念')]

    sql = f"SELECT FirstIndustryName, FirstIndustryCode FROM AStockIndustryDB.LC_ExgIndustry WHERE FirstIndustryName = '{value}' LIMIT 1"
    res = exec_sql_s(sql)
    if res:
        return [(res, f'{value}是AStockIndustryDB.LC_ExgIndustry的一级行业名称')]
    else:
        sql = f"SELECT SecondIndustryName, SecondIndustryCode FROM AStockIndustryDB.LC_ExgIndustry WHERE SecondIndustryName = '{value}' LIMIT 1"
        res = exec_sql_s(sql)
        if res:
            return [(res, f'{value}是AStockIndustryDB.LC_ExgIndustry的二级行业名称')]
        else:
            sql = f"SELECT ThirdIndustryName, ThirdIndustryCode FROM AStockIndustryDB.LC_ExgIndustry WHERE ThirdIndustryName = '{value}' LIMIT 1"
            res = exec_sql_s(sql)
            if res:
                return [(res, f'{value}是AStockIndustryDB.LC_ExgIndustry的三级行业名称')]
            else:
                sql = f"SELECT FourthIndustryName, FourthIndustryCode FROM AStockIndustryDB.LC_ExgIndustry WHERE FourthIndustryName = '{value}' LIMIT 1"
                res = exec_sql_s(sql)
                if res:
                    return [(res, f'{value}是AStockIndustryDB.LC_ExgIndustry的四级行业名称')]
    return []



def _process_jj(value):
    from fuzzywuzzy import fuzz
    prompt = f'你会为基金名称找到一个搜索用的关键词，一般是公司名+产品名，示例：嘉实致元42个月定期债券基金->嘉实致元,仅仅只输出一个单词，基金名称为：{value} 关键词:'
    key_word = llm(prompt)
    sql = f"SELECT ChiName,CompanyCode,ChiName,InnerCode FROM ConstantDB.SecuMain WHERE ChiName LIKE '%{key_word}%'"
    score = 0
    res=None
    res_li = exec_sql_s(sql)
    for i in res_li:
        if fuzz.ratio(i['ChiName'],value) > score:
            score = fuzz.ratio(i['ChiName'],value)
            res = i


    if res:
        res['警告'] = '用户输入错误，请依据预处理中的名称'
        return [(res, 'InstitutionDB.LC_InstiArchive')]
    else:

        return []


def process_jj(value):
    sql = f"SELECT ChiName, CompanyCode FROM InstitutionDB.LC_InstiArchive WHERE ChiName LIKE '%{value}%'"
    res = exec_sql_s(sql)
    if res:
        return [(res, 'InstitutionDB.LC_InstiArchive')]
    else:
        return []


def process_items(item_list):
    res_list = []
    for item in item_list:
        key, value = list(item.items())[0]
        if key == "基金名称":
            res_list.extend(process_company_name(value,is_jj=True))
        elif key == "上市公司名称":
            res_list.extend(process_company_name(value))
        elif key == "基金公司简称":
            res_list.extend(process_jj(value))
        elif key == "代码":
            res_list.extend(process_code(value))
        elif key == "板块名称":
            res_list.extend(process_bk(value))
        else:
            print(f"无法识别的键：{key}")
    res_list = [i for i in res_list if i]
    res = ''
    tables = []
    for i, j in res_list:
        tables.append(j)
        res += f"预处理程序通过表格：{j} 查询到以下内容：\n {json.dumps(i, ensure_ascii=False, indent=1)} \n"
    return res, tables


def process_question(question):
    prompt = '''
你将会进行命名实体识别任务，并输出实体json

你只需要识别以下4种实体：
-上市公司名称
-代码
-基金名称
-基金公司简称
-板块名称

只有出现了才识别，问题不需要回答

其中，公司名称可以是全称，简称，拼音缩写，对于不规范的公司全称，只识别简称部分。
代码包含股票代码和基金代码
基金名称包含债券型基金，


以下是几个示例：
user:实体识别任务：```唐山港集团股份有限公司是什么时间上市的（回答XXXX-XX-XX）
当年一共上市了多少家企业？
这些企业有多少是在北京注册的？```
assistant:```json
[{"上市公司名称":"唐山港集团股份有限公司"}]
```
说明:唐山港集团股份有限公司是标准的公司名

user:实体识别任务：```利亚德光电在2019到2021年这三年间技术人员和生产人员占职工总数的比例分别是多少？```
assistant:```json
[{"上市公司名称":"利亚德"}]
```
说明:利亚德光电不是标准的公司名，只识别`利亚德`

user:实体识别任务：```JD的职工总数有多少人？
该公司披露的硕士或研究生学历（及以上）的有多少人？
20201月1日至年底退休了多少人？```
assistant:```json
[{"上市公司名称":"JD"}]
```
user:实体识别任务：```RXW中文全称是什么？```
assistant:```json
[{"上市公司名称":"RXW"}]
```
user:实体识别任务：```600872的全称、A股简称、法人、法律顾问、会计师事务所及董秘是？
该公司实控人是否发生改变？如果发生变化，什么时候变成了谁？是哪国人？是否有永久境外居留权？（回答时间用XXXX-XX-XX）```
assistant:```json
[{"代码":"600872"}]
```
user:实体识别任务：```华夏鼎康债券A在2019年的分红次数是多少？每次分红的派现比例是多少？
基于上述分红数据，在2019年最后一次分红时，如果一位投资者持有1000份该基金，税后可以获得多少分红收益？```
assistant:```json
[{"基金名称":"华夏鼎康债券A"}]
```
user:实体识别任务：```易方达基金管理有限公司在19年成立了多少支基金？
哪支基金的规模最大？
这支基金20年最后一次分红派现比例多少钱？```
assistant:```json
[{"基金公司简称":"易方达"}]
```
user:实体识别任务：```化工纳入过多少个子类概念？```
assistant:```json
[{"板块名称":"化工"}]
```
user:实体识别任务：```今天是2021年12月24日，创近半年新高的股票有几只？```
assistant:```json
[]
```

'''

    messages = [{'role': 'system', 'content': prompt},

                {'role': 'user', 'content': f'''实体识别任务：```{question}```'''}]

    a1 = llm(messages)
    a2 = super_eval(a1)
    return process_items(a2)


if __name__ == '__main__':
    question = 'Q1:A股-在2021年全年，厦门钨业的最高收盘价是多少，出现在哪一天？'
    res = process_question(question)

    print(res)
